#include "libyuv_ext.h"

#ifdef __cplusplus
extern "C" {
#endif

int ExceptionCheck(JNIEnv *env) {
    if (env->ExceptionCheck()) {
        LOGE("JNI Exception is happened.");
        env->ExceptionDescribe();// writes to logcat
        env->ExceptionClear();
        return 1;
    }
    return 0;
}

JNIEXPORT void JNICALL
Java_com_andova_libyuv_LibYUV_00024Companion_YUV422ToNV21(
        JNIEnv *env, jobject,
        jobject src_y, jobject src_u,
        jobject src_v, jbyteArray dst_yuv,
        jint angle, jint width, jint height) {
    if (ExceptionCheck(env))return;
    LOGI("TestCpuFlag:%d", libyuv::TestCpuFlag(libyuv::kCpuHasNEON));
    auto *ybuf = (uint8_t *) env->GetDirectBufferAddress(src_y);
    auto *ubuf = (uint8_t *) env->GetDirectBufferAddress(src_u);
    auto *vbuf = (uint8_t *) env->GetDirectBufferAddress(src_v);
    if (!ybuf || !ubuf || !vbuf)return;
    jbyte *yuvbuf = env->GetByteArrayElements(dst_yuv, nullptr);

#if false
    // Allocate y and u and v buffers
    align_buffer_64(plane_y, width * height + SUBSAMPLE(width, 2) * SUBSAMPLE(height, 2) * 2);
    uint8_t *plane_u = plane_y + width * height;
    uint8_t *plane_v = plane_u + SUBSAMPLE(width, 2) * SUBSAMPLE(height, 2);

#if false
    I420Copy((const uint8_t *) ybuf, width,
             (const uint8_t *) ubuf, SUBSAMPLE(width, 2),
             (const uint8_t *) vbuf, SUBSAMPLE(width, 2),
             plane_y, width,
             plane_u, SUBSAMPLE(width, 2),
             plane_v, SUBSAMPLE(width, 2),
             width, height);
#endif

    int halfwidth = SUBSAMPLE(width, 2);
    int halfheight = SUBSAMPLE(height, 2);
    RotatePlane90(ybuf, width, plane_y, height, width, height);
    RotatePlane90(ubuf, halfwidth, plane_u, halfheight, halfwidth, halfheight);
    RotatePlane90(vbuf, halfwidth, plane_v, halfheight, halfwidth, halfheight);
    libyuv::I422ToNV21(plane_y, height,
                       plane_u, halfheight,
                       plane_v, halfheight,
                       (uint8_t *) yuvbuf, height,
                       (uint8_t *) yuvbuf + width * height, height,
                       height, width);

#if false
    libyuv::I420Rotate((const uint8_t *) ybuf, width,
                       (const uint8_t *) ubuf, SUBSAMPLE(width, 2),
                       (const uint8_t *) vbuf, SUBSAMPLE(width, 2),
                       plane_y, height,
                       plane_u, SUBSAMPLE(height, 2),
                       plane_v, SUBSAMPLE(height, 2),
                       width, height, libyuv::kRotate90);
#endif

    free_aligned_buffer_64(plane_y);
#elif false
    align_buffer_64(plane_y, width * height + SUBSAMPLE(width, 2) * SUBSAMPLE(height, 2) * 2);
    libyuv::I422ToNV21((uint8_t *) ybuf, width,
                       (uint8_t *) ubuf, SUBSAMPLE(width, 2),
                       (uint8_t *) vbuf, SUBSAMPLE(width, 2),
                       (uint8_t *) yuvbuf, width,
                       (uint8_t *) yuvbuf + width * height, width,
                       width, height);
    free_aligned_buffer_64(plane_y);
#else
    switch (angle) {
        case libyuv::kRotate90:
        case libyuv::kRotate180:
        case libyuv::kRotate270: {
            align_buffer_64(plane_yuv, env->GetArrayLength(dst_yuv));
            YUV422ToNV21(ybuf, width,
                         ubuf, SUBSAMPLE(width, 2),
                         vbuf, SUBSAMPLE(width, 2),
                         plane_yuv, width,
                         plane_yuv + width * height, width,
                         width, height);
            RotateNV21(plane_yuv, (uint8_t *) yuvbuf, angle, width, height);
            free_aligned_buffer_64(plane_yuv);
            break;
        }
        default: {
            int y_len = width * height;
            auto *dst_y = (uint8_t *) yuvbuf;
            uint8_t *dst_vu = dst_y + y_len;
            YUV422ToNV21(ybuf, width,
                         ubuf, SUBSAMPLE(width, 2),
                         vbuf, SUBSAMPLE(width, 2),
                         dst_y, width,
                         dst_vu, width,
                         width, height);
        }
    }
#endif

    env->ReleaseByteArrayElements(dst_yuv, yuvbuf, 0);
}

void YUV422ToNV21(const uint8_t *src_y, int src_stride_y,
                  const uint8_t *src_u, int src_stride_u,
                  const uint8_t *src_v, int src_stride_v,
                  uint8_t *dst_y, int dst_stride_y,
                  uint8_t *dst_vu, int dst_stride_vu,
                  int width, int height) {
    int halfwidth = SUBSAMPLE(width, 2);
    int halfheight = SUBSAMPLE(height, 2);
    if (!src_y || !src_u || !src_v || !dst_y || !dst_vu || width <= 0 || height == 0) return;
    // Negative height means invert the image.
    if (height < 0) {
        height = -height;
        halfheight = (height + 1) >> 1;
        src_y = src_y + (height - 1) * src_stride_y;
        src_u = src_u + (halfheight - 1) * src_stride_u;
        src_v = src_v + (halfheight - 1) * src_stride_v;
        src_stride_y = -src_stride_y;
        src_stride_u = -src_stride_u;
        src_stride_v = -src_stride_v;
    }

    // Negative height means invert the image.
    if (height < 0) {
        height = -height;
        dst_y = dst_y + (height - 1) * dst_stride_y;
        dst_stride_y = -dst_stride_y;
    }
    // Coalesce rows.
    if (src_stride_y == width && dst_stride_y == width) {
        width *= height;
        height = 1;
        src_stride_y = dst_stride_y = 0;
    }
    // Nothing to do.
    if (src_y == dst_y && src_stride_y == dst_stride_y) {
        return;
    }

    // Negative height means invert the image.
    if (halfheight < 0) {
        halfheight = -halfheight;
        dst_vu = dst_vu + (halfheight - 1) * dst_stride_vu;
        dst_stride_vu = -dst_stride_vu;
    }
    // Coalesce rows.
    if (src_stride_u == halfwidth && src_stride_v == halfwidth &&
        dst_stride_vu == halfwidth * 2) {
        halfwidth *= halfheight;
        halfheight = 1;
        src_stride_u = src_stride_v = dst_stride_vu = 0;
    }

    libyuv::CopyPlane(src_y, width, dst_y, width, width, height);
#if false
    libyuv::HalfMergeUVPlane(plane_u, src_stride_u, plane_v, src_stride_v, dst_vu, dst_stride_vu, halfwidth, halfheight);
#endif

#if true
    if (!IS_ALIGNED(width, 16)) {
        LOGE("is not aligned 16");
        return;
    }
    int stride_uv = halfwidth * 2; // 经测试，发现subs每8位一减，很容易崩溃，故换成16，但依然有问题
    for (int y = 0; y < halfheight; ++y) {
        asm volatile(
        "MergeI422VU:                   \n"
        "vld1.8         q0, [%0]!       \n" // load 16 V
        "vld1.8         q1, [%1]!       \n" // load 16 U
        "vshl.u16       q0, q0, #8      \n" // <<
        "vqrshrn.u16    d0, q0, #8      \n" // half size and >>
        "vshl.u16       q1, q1, #8      \n" // <<
        "vqrshrn.u16    d1, q1, #8      \n" // half size and >>
        "subs           %3, %3, #16     \n" // 16 processed per loop
        "vst2.8         {d0, d1}, [%2]! \n" // store 8 pairs of VU
        "bgt            MergeI422VU     \n"
        : "+r"(src_v),                      // %0
        "+r"(src_u),                        // %1
        "+r"(dst_vu),                       // %2
        "+r"(stride_uv)                     // %3  // Output registers
        :                                   // Input registers
        : "cc", "memory", "q0", "q1"        // Clobber List
        );
        src_u += src_stride_u;
        src_v += src_stride_v;
        dst_vu += dst_stride_vu;
    }
#else
    int len = halfwidth * halfheight;
    align_buffer_64(plane_u, len);
    align_buffer_64(plane_v, len);
    for (int i = 0; i < len; ++i) {
        plane_u[i] = src_u[i * 2];
        plane_v[i] = src_v[i * 2];
    }
    libyuv::MergeUVPlane(plane_v, halfwidth, plane_u, halfwidth, dst_vu, width, halfwidth, halfheight);
    free_aligned_buffer_64(plane_u);
    free_aligned_buffer_64(plane_v);
#endif
}

void RotateNV21(uint8_t *src, uint8_t *dst, jint angle, jint width, jint height) {
    int halfwidth = SUBSAMPLE(width, 2);
    int halfheight = SUBSAMPLE(height, 2);
    uint8_t *src_vu = src + width * height;
    align_buffer_64(plane_u, halfwidth * halfheight);
    align_buffer_64(plane_v, halfwidth * halfheight);
    switch (angle) {
        case libyuv::kRotate0:
            libyuv::CopyPlane(src, width, dst, width, width, height);
            libyuv::CopyPlane(src_vu, width, dst + width * height, width, width, halfheight);
            break;
        case libyuv::kRotate90:
            libyuv::RotatePlane90(src, width, dst, height, width, height);
            libyuv::RotateUV90(src_vu, width, plane_v, halfheight, plane_u, halfheight, halfwidth, halfheight);
            libyuv::MergeUVPlane(plane_v, halfheight, plane_u, halfheight, dst + width * height, height, halfheight, halfwidth);
            break;
        case libyuv::kRotate180:
            libyuv::RotatePlane180(src, width, dst, width, width, height);
            libyuv::RotateUV180(src_vu, width, plane_v, halfwidth, plane_u, halfwidth, halfwidth, halfheight);
            libyuv::MergeUVPlane(plane_v, halfwidth, plane_u, halfwidth, dst + width * height, width, halfwidth, halfheight);
            break;
        case libyuv::kRotate270:
            libyuv::RotatePlane270(src, width, dst, height, width, height);
            libyuv::RotateUV270(src_vu, width, plane_v, halfheight, plane_u, halfheight, halfwidth, halfheight);
            libyuv::MergeUVPlane(plane_v, halfheight, plane_u, halfheight, dst + width * height, height, halfheight, halfwidth);
            break;
        default:
            LOGE("error angle value %d", angle);
            break;
    }
    free_aligned_buffer_64(plane_u);
    free_aligned_buffer_64(plane_v);
}

void RotateYUV(const uint8_t *src_y, uint8_t *dst_y,
               const uint8_t *src_u, uint8_t *dst_u,
               const uint8_t *src_v, uint8_t *dst_v,
               int angle, int width, int height) {
    int halfwidth = SUBSAMPLE(width, 2);
    int halfheight = SUBSAMPLE(height, 2);
    switch (angle) {
        case libyuv::kRotate0:
            libyuv::CopyPlane(src_y, width, dst_y, width, width, height);
            libyuv::CopyPlane(src_u, halfwidth, dst_u, halfwidth, halfwidth, halfheight);
            libyuv::CopyPlane(src_v, halfwidth, dst_v, halfwidth, halfwidth, halfheight);
            break;
        case libyuv::kRotate90:
            libyuv::RotatePlane90(src_y, width, dst_y, height, width, height);
            libyuv::RotatePlane90(src_u, halfwidth, dst_u, halfheight, halfwidth, halfheight);
            libyuv::RotatePlane90(src_v, halfwidth, dst_v, halfheight, halfwidth, halfheight);
            break;
        case libyuv::kRotate180:
            libyuv::RotatePlane180(src_y, width, dst_y, width, width, height);
            libyuv::RotatePlane180(src_u, halfwidth, dst_u, halfwidth, halfwidth, halfheight);
            libyuv::RotatePlane180(src_v, halfwidth, dst_v, halfwidth, halfwidth, halfheight);
            break;
        case libyuv::kRotate270:
            libyuv::RotatePlane270(src_y, width, dst_y, height, width, height);
            libyuv::RotatePlane270(src_u, halfwidth, dst_u, halfheight, halfwidth, halfheight);
            libyuv::RotatePlane270(src_v, halfwidth, dst_v, halfheight, halfwidth, halfheight);
            break;
        default:
            LOGE("error angle value %d", angle);
            break;
    }
}

[[deprecated("Inefficient")]]void YUV420Rotate90(uint8_t *dst, const uint8_t *src, int src_width, int src_height) {
    static int nWidth = 0, nHeight = 0;
    static int wh = 0;
    static int uvHeight = 0;
    if (src_width != nWidth || src_height != nHeight) {
        nWidth = src_width;
        nHeight = src_height;
        wh = src_width * src_height;
        uvHeight = src_height >> 1;
    }

    //旋转Y
    int k = 0;
    for (int i = 0; i < src_width; i++) {
        int nPos = 0;
        for (int j = 0; j < src_height; j++) {
            dst[k] = src[nPos + i];
            k++;
            nPos += src_width;
        }
    }

    for (int i = 0; i < src_width; i += 2) {
        int nPos = wh;
        for (int j = 0; j < uvHeight; j++) {
            dst[k] = src[nPos + i];
            dst[k + 1] = src[nPos + i + 1];
            k += 2;
            nPos += src_width;
        }
    }
}

JNIEXPORT jbyteArray JNICALL
Java_com_andova_libyuv_LibYUV_00024Companion_GetByteArray(JNIEnv *env, jobject, jobject src, jint len) {
    auto *src_buf = (jbyte *) env->GetDirectBufferAddress(src);
    if (!src_buf)return nullptr;
    jbyteArray dst = env->NewByteArray(len);
    env->SetByteArrayRegion(dst, 0, len, src_buf);
    return dst;
}

JNIEXPORT void JNICALL
Java_com_andova_libyuv_LibYUV_00024Companion_RotateNV21(JNIEnv *env, jobject,
                                                        jbyteArray src, jbyteArray dst,
                                                        jint angle, jint width, jint height) {
    jbyte *src_buf = env->GetByteArrayElements(src, nullptr);
    jbyte *dst_buf = env->GetByteArrayElements(dst, nullptr);
    RotateNV21((uint8_t *) src_buf, (uint8_t *) dst_buf, angle, width, height);
    env->ReleaseByteArrayElements(src, src_buf, 0);
    env->ReleaseByteArrayElements(dst, dst_buf, 0);
}

#ifdef __cplusplus
}
#endif